#include <vector>
#include <cstring>
#include <algorithm>

#include "painter.hh"
#include "debug.hh"

namespace icesp
{
    painter::painter()
        : colors(nullptr), node_number_count(nullptr), label(nullptr), unlabel(nullptr)
    {
    }

    painter::~painter()
    {
        release();
    }

    int painter::bfs_delta_depth_coloring(graph::adj_matrix & g, int source, int delta_depth)
    {
        n = g.n;
        release_and_allocate();

        auto & edges = g.edges;
        auto & edge_index = g.edge_index;

        // TODO
        auto const unvisited = -1;
        memset(colors, 0xff, sizeof(int) * n);
        std::vector<int> depth(n);
        std::vector<int> q{source};
        auto & father = colors;
        std::vector<int>::size_type it{0};
        while (it != q.size()) {
            auto const from = q[it];
            for (auto i = edge_index[from]; i < edge_index[from + 1]; i++) {
                auto const to = edges[i].to;
                if (father[to] == unvisited) {
                    father[to] = from;
                    depth[to] = depth[from] + 1;
                    q.push_back(to);

                }
            }
            ++it;
        }

        father[source] = unvisited;
        auto current_thread_id = 0;
        for (auto i : q) {
            if ((depth[i] & (delta_depth - 1)) == 0) {
                colors[i] = current_thread_id;
                current_thread_id++;
            } else
                colors[i] = colors[father[i]];
        }
        return current_thread_id;
    }

    void painter::refine(int nthread)
    {
        this->nthread = nthread;

        release_node_number_count();
        node_number_count = new int[nthread + n + n];
        std::memset(node_number_count, 0, sizeof(int) * nthread);

        for (auto i = 0; i < n; i++) {
            colors[i] %= nthread;
            node_number_count[colors[i]]++;
        }

          label = node_number_count + nthread;
        unlabel = node_number_count + nthread + n;
        int* tmp = new int[nthread];
        for (auto i = 0, cur = 0; i < nthread; i++) {
            tmp[i] = cur;
            cur += node_number_count[i];
        }
        for (auto i = 0; i < n; i++) {
            label[i] = tmp[colors[i]];
            unlabel[label[i]] = i;
            tmp[colors[i]]++;
        }
        delete[] tmp;
    }

    void painter::from_array(int len, int* color_arr)
    {
        n = len;
        release_and_allocate();
        std::copy(color_arr, color_arr + n, colors);
    }

    int const* painter::get_colors  () const { return colors;  }
    int const* painter::get_label   () const { return label;   }
    int const* painter::get_unlabel () const { return unlabel; }
    int const* painter::get_node_number_count () const { return node_number_count; }

    void painter::release()
    {
        if (colors != nullptr) {
            delete[] colors;
            colors = nullptr;
        }
        release_node_number_count();
    }

    void painter::release_node_number_count()
    {
        if (node_number_count != nullptr) {
            delete[] node_number_count;
            node_number_count = label = unlabel = nullptr;
        }
    }

    void painter::release_and_allocate()
    {
        release();
        colors = new int[n];
    }
}

